import cv2
import numpy as np
from torch.utils.data import Dataset
from torchvision import transforms
import urllib.request


class OSSDataset(Dataset):
    """
    Dataset for test prediction
    """

    def __init__(self, items, height, width, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
        """
        items:[
                {'id': 1,'url':'ww.cdscf.dscd'},
                {'id': 2,'url':'ww.cdscf.csdcs'},
             ]
        """
        self.items = items
        self.num_samples = len(items)
        self.height = height
        self.weight = width
        self.transform = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Resize((height, width)),
                transforms.Normalize(mean=mean, std=std)
            ]
        )

    def __getitem__(self, idx):
        id, url = self.items[idx]['id'], self.items[idx]['url']
        resp = urllib.request.urlopen(url)
        image_data = np.asarray(bytearray(resp.read()), dtype=np.uint8)
        image = cv2.cvtColor(cv2.imdecode(image_data, cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB)
        image = self.transform(image)
        return id, image

    def __len__(self):
        return self.num_samples


class UnsupervisedDataset(Dataset):
    def __init__(self, image_paths, height, width, mean=[0.485, 0.456, 0.406],
                 std=[0.229, 0.224, 0.225]):
        super().__init__()
        self.image_paths = image_paths
        self.height = height
        self.width = width
        self.mean = mean
        self.std = std
        self.transform_image = transforms.Compose([
            transforms.ToTensor(),
            transforms.Resize((height, width)),
            transforms.Normalize(self.mean, self.std)
        ])

    def __getitem__(self, index):
        image_path = self.image_paths[index]
        img = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
        img = self.transform_image(img)
        return img

    def __len__(self):
        return len(self.image_paths)


if __name__ == '__main__':
    items = [
        {'id': 1, 'url': 'https://msjava.oss-cn-qingdao.aliyuncs.com/temp/00b7fb703.jpg'},
        {'id': 2, 'url': 'https://msjava.oss-cn-qingdao.aliyuncs.com/temp/00bbcd9af.jpg'},
        {'id': 3, 'url': 'https://msjava.oss-cn-qingdao.aliyuncs.com/temp/1bbf4b4c0.jpg'},
    ]
    dataset = OSSDataset(items, 256, 1600)
    print(dataset.__getitem__(0)[1].shape)
